#coding=utf-8
import json
import pymongo
from bson import json_util
from Core.DataSeries import *
import Core.Gadget as Gadget

MAX_DATABASE_SCAN = 5000000

class MongoDB(object):

    def __init__(self, address, port, password = None):

        print("Connecting Database")
        self.address = address
        self.port = port
        self.username = "admin"
        self.password = password
        # Create a Connection "mongodb://mongodb0.example.net:27019"
        self.client = pymongo.MongoClient("mongodb://" + address + ":" + port)
        if self.password != None:
            db = self.client["admin"]
            db.authenticate("admin","iKidzyueya2016",mechanism='SCRAM-SHA-1')
        print("Connected Database")

    def addItem(self, databaseName, collectionName, target, items):
        db = self.client[databaseName]
        coll = db[collectionName]
        update = {}
        update["$set"] = items  # {"Values.test3":3}
        coll.update_one(target, update, True)  # {"Key":"IC.CFE"}

    def addListItem(self, databaseName, collectionName, target, items):
        db = self.client[databaseName]
        coll = db[collectionName]
        update = {}
        update["$push"] = items  # {"Values":3}
        coll.update_one(target, update, True)  # {"Key":"IC.CFE"}

    def creatIndex(self, databaseName, collectionName, indexKeyName):
        db = self.client[databaseName]
        coll = db[collectionName]
        coll.create_index(indexKeyName)
        print("Created Index: " + indexKeyName + " for " + collectionName)

    def close(self):
        self.client.close()

    def delete(self, databaseName, collectionName, query):
        db = self.client[databaseName]
        coll = db[collectionName]
        coll.delete_many(query)
        # result = db.restaurants.delete_many({"borough": "Manhattan"})

    def DeleteCollections(self, databaseName, pattern = ""):
        db = self.client[databaseName]
        names = db.collection_names(False)
        for name in names:
            if pattern in name:
                coll = db[name]
                coll.drop()
                print("Delete Collection:" + name)

    def DeleteCollection(self, databaseName, collectionName):
        db = self.client[databaseName]
        names = db.collection_names(False)
        if collectionName in names:
            coll = db[collectionName]
            coll.drop()
            print("Delete Collection:" + collectionName)


    def deleteDataSeries(self, dataSeriesName):
        symbol = self.parseSymbol(dataSeriesName)
        type = self.parseDataType(dataSeriesName)
        name = self.parseDataName(dataSeriesName)

        if type == "Factor":
            self.delete("Factor", name, {"Symbol": symbol})

    def findWithFilter(self, databaseName, collectionName, filter={}, sort=None, projection={}):

        db = self.client[databaseName]
        #if self.password != None:
        #    db.authenticate(self.username, self.password)

        skip = 0
        limit = 0
        if "skip" in filter:
            skip = filter["skip"]
            filter.pop("skip")
        if "limit" in filter:
            limit = filter["limit"]
            filter.pop("limit")

        #
        if len(projection) !=0:
            cursor = db[collectionName].find(filter,skip=skip,limit=limit,sort=sort,projection=projection).max_scan(MAX_DATABASE_SCAN)
        else:
            cursor = db[collectionName].find(filter, skip=skip, limit=limit, sort=sort).max_scan(MAX_DATABASE_SCAN)
        objList = []

        #
        for document in cursor:
            #print(document)
            jsonDoc = json.dumps(document, default=json_util.default)
            jsonObj = json.loads(jsonDoc, object_hook=json_util.object_hook)
            #if "_id" in jsonObj:
            #    del jsonObj["_id"]
            objList.append(jsonObj)
        return objList

    def find(self, databaseName, collectionName, beginDateTime=None, endDateTime=None, query={}, sort=None, asc=True, projection={}):
        #
        db = self.client[databaseName]
        query["StdDateTime"] = {}
        if beginDateTime != None:
            query["StdDateTime"]["$gte"] = beginDateTime
        if endDateTime != None:
            query["StdDateTime"]["$lte"] = endDateTime

        if not query["StdDateTime"]:
            del query["StdDateTime"]

        sortDoc = []
        if sort != None:
            sortDoc = sort
        else:
            if asc:
                sort = pymongo.ASCENDING
            else:
                sort = pymongo.DESCENDING
            sortDoc = [("StdDateTime", sort)]

        objList = self.findWithFilter(databaseName=databaseName,
                                      collectionName=collectionName,
                                      filter=query,
                                      sort=sortDoc,
                                      projection=projection)
        #
        return objList

    def findAll(self,databaseName, collectionName):
        db = self.client[databaseName]
        cursor = db[collectionName].find().max_scan(100000)
        objList = []
        for document in cursor:
            #print(document)
            jsonDoc = json.dumps(document, default=json_util.default)
            jsonObj = json.loads(jsonDoc, object_hook=json_util.object_hook)
            objList.append(jsonObj)
        return objList

    def getCollection(self, databaseName, collectionName):
        db = self.client[databaseName]
        coll =  db[collectionName]
        return coll

    def GetCollectionNames(self, databaseName, pattern):
        db = self.client[databaseName]
        names = db.collection_names(False)
        filteredNames = []
        for name in names:
            if pattern in name:
                filteredNames.append(name)
        return filteredNames

    def getDataSeries(self, dataSeriesName, datetime1=None, datetime2=None, instrumentType=None, asObject=True):

        #
        type = self.parseDataType(dataSeriesName)
        name = self.parseDataName(dataSeriesName)
        symbol = self.parseSymbol(dataSeriesName)
        #
        sort = [("StdDateTime", pymongo.ASCENDING)]
        query = {}

        #
        query["StdDateTime"] = {}
        if datetime1 != None:
            query["StdDateTime"]["$gte"] = datetime1
        if datetime2 != None:
            query["StdDateTime"]["$lte"] = datetime2
        if not query["StdDateTime"]:
            del query["StdDateTime"]

        # collectionName = name + "_Time_86400_Bar"
        if type == "Bar":
            collectionName = dataSeriesName
            # if instrumentType == None:
            #    # stockList = self.find("Instruments","Stock",fields={"Symbol":1})
            #    print("Unknown Bar Type")
            #    return
            if instrumentType == "Stock":  # and ( size == 86400 or size == 60):
                databaseName = "Quote"
            elif instrumentType == "Index":
                databaseName = "Index"
            elif instrumentType == "Future":
                databaseName = "Future"
                # query["DataSeries"] = dataSeriesName
            elif instrumentType == "Bond":
                databaseName = "Bond"
                # query["DataSeries"] = dataSeriesName
            elif instrumentType == "MutualFund":
                databaseName = "MutualFund"
            else:
                # ---default stock---
                databaseName = "Quote"
                # print("Unknown Bar Type")
                # return

        elif type == "Fundamental":
            databaseName = "Fundamental"
            collectionName = dataSeriesName

        elif type == "Portfolio":
            databaseName = "Portfolio"
            collectionName = "PortfolioHistData"

        elif type == "InstrumentList":
            databaseName = "Instruments"
            collectionName = "InstrumentList"
            query["Symbol"] = symbol

        elif type == "Factor":
            databaseName = "Factor"
            collectionName = name
            query["Symbol"] = symbol

        elif type == "Label":  # ?
            print("Load Label DataSeries")

        elif type == "Event":  # ?
            print("Load Event DataSeries")

        else:
            print("Unknown DataSerie Type: " + type)
            return

        #
        if asObject:
            # ---
            db = self.client[databaseName]
            # cursor = db[collectionName].find({"StdDateTime": {"$gt": beginDateTime, "$lt": endDateTime} }).max_scan(100000)
            cursor = db[collectionName].find(query).sort(sort)
            # ---
            dataSeries = DataSeries(dataSeriesName)
            for document in cursor:
                # print(document)
                jsonDoc = json.dumps(document, default=json_util.default)
                jsonObj = json.loads(jsonDoc, object_hook=json_util.object_hook)
                del jsonObj["_id"]
                dataSeries.Add(jsonObj)
        else:
            dataSeries = self.findWithFilter(databaseName, collectionName, query, sort)

        return dataSeries

    def increase(self, databaseName, collectionName, target, dataObject):
        db = self.client[databaseName]
        coll =  db[collectionName]
        update = {}
        update["$inc"] = dataObject
        coll.update_one(target, update, False)

    def insert(self, databaseName, collectionName, dataObject):
        db = self.client[databaseName]
        coll = db[collectionName]
        coll.insert(dataObject)
        #coll.insert_one()

    def insert_many(self, databaseName, collectionName, dataList):
        db = self.client[databaseName]
        coll = db[collectionName]
        coll.insert_many(dataList)

    def insertIfNotExisted(self, databaseName, collectionName, dataObject, target):
        db = self.client[databaseName]
        coll =  db[collectionName]
        doc = coll.find_one(target)
        if doc == None:
            coll.insert(dataObject)
            return True
        return False

    def parseSymbol(self, dataSeriesName):
        names = dataSeriesName.split('_')
        return names[0]

    def parseDataType(self, dataSeriesName):
        names = dataSeriesName.split('_')
        return names[len(names)-1]

    def parseDataName(self, dataSeriesName):
        names = dataSeriesName.split('_')
        return names[len(names)-2]

    def parseSize(self, dataSeriesName):
        names = dataSeriesName.split('_')
        if len(names) > 2:
            return int(names[2])
        return 0

    def saveDataSeries(self, dataSeries, instrumentType=None, isUpSert=True):

        dataSeriesName = dataSeries.name
        print(str(datetime.now()) + " Saving DataSeries " + dataSeriesName)
        #
        if dataSeries.Count() == 0:
            print(dataSeriesName + " saveDataSeries::No Data to Save")
            return

        symbol = self.parseSymbol(dataSeriesName)
        type = self.parseDataType(dataSeriesName)
        name = self.parseDataName(dataSeriesName)

        databaseName = ""
        collectionName = ""
        keys = []
        insertType = 0
        if type == "Bar":
            size = self.parseSize(dataSeriesName)
            dataBaseStructureType = "BySymbol"
            if instrumentType == None:
                #stockList = self.find("Instruments","Stock",fields={"Symbol":1})
                print("Unknown Bar Type")
                return
            elif (instrumentType == "Stock") and (size == 86400 or size == 60 or size == 300):
                databaseName = "Quote"
                collectionName = dataSeriesName
                self.saveCollection(dataSeries.items, databaseName, collectionName)
            elif instrumentType == "Index":
                databaseName = "Index"
                collectionName = dataSeriesName
                self.saveCollection(dataSeries.items, databaseName, collectionName)
            elif instrumentType == "Future":
                databaseName = "Future"
                collectionName = dataSeriesName
                self.saveCollection(dataSeries.items, databaseName, collectionName)
            elif instrumentType == "Bond":
                databaseName = "Bond"
                # query["DataSeries"] = dataSeriesName
            elif instrumentType == "MutualFund":
                collectionName = dataSeriesName
                self.saveCollection(dataSeries.items, databaseName, collectionName)
            else:
                print("Unknown Bar Type")
                return

        elif type == "Fundamental":
            databaseName = "Fundamental"
            collectionName = dataSeriesName
            self.saveCollection(dataSeries.items, databaseName, collectionName)

        elif type == "InstrumentList":
            databaseName = "Instruments"
            collectionName = "InstrumentList"

        elif type == "Factor":
            databaseName = "Factor"
            collectionName = name
            self.saveCollection(dataSeries.items, databaseName, collectionName, keys=["Symbol"], isUpSert=isUpSert)

        elif type == "Label":# ?
            print("Save Label DataSeries")

        elif type == "Event":# ?
            print("Save Event DataSeries")
        #
        print(str(datetime.now()) + " Saved DataSeries " + dataSeriesName + " # of Elements " + str(dataSeries.Count()))

    def saveCollection(self, dataList, databaseName, collectionName, keys=[], isUpSert=True):

        # Loop to Insert DataBase
        # print(str(datetime.now()) + " Saving DataSeries " + collectionName)

        # ---Determine if the first time to save data (if the dataSeries exist?)---
        isFirstSaveData = False
        collectionNames = self.client[databaseName].collection_names(False)
        if collectionName not in collectionNames:
            isFirstSaveData = True
            isUpSert = False # First Save data, unnecessary to use Upsert Command

        # ---Create Key for each Element---
        for data in dataList:
            if "Key" not in data:
                localDateTime = Gadget.ToLocalDateTime(data["StdDateTime"])
                key = data["DataSeries"] + "_" + Gadget.ToDateTimeString(localDateTime) # Name + DateTime as identifier
                data["Key"] = key

        # if first to create DatSeries, Build Index
        if isFirstSaveData:
            db = self.client[databaseName]
            coll = db[collectionName]
            coll.create_index("Key")
            coll.create_index("StdDateTime")
            for key in keys:
                coll.create_index(key)
            #
            #self.insert_many(databaseName, collectionName, dataList)
            self.upsert_many(databaseName, collectionName, ["Key"], dataList)
        else:
            # for data in dataList:
            #    target = {"Key": data["Key"]}
            #    self.upsert(databaseName, collectionName, target, data)
            if isUpSert:
                self.upsert_many(databaseName, collectionName, ["Key"], dataList)
            else:
                self.insert_many(databaseName, collectionName, dataList)

        print(str(datetime.now()) + " Saved DataSeries " + collectionName + " # of Elements " + str(len(dataList)))

    def Stats(self, databaseName, collectionName):
        db = self.client[databaseName]
        ret = db.command("collstats", collectionName)
        return ret

    def upsert(self, databaseName, collectionName, target, dataObject):
        if "_id" in dataObject:
            del dataObject["_id"]
        #
        db = self.client[databaseName]
        coll = db[collectionName]
        update = {}
        update["$set"] = dataObject
        coll.update_one(target, update, True)

    def upsert_many(self, databaseName, collectionName, keys, dataList):
        db = self.client[databaseName]
        coll = db[collectionName]
        requests = []
        for data in dataList:
            filter = {}
            for key in keys:
                filter[key] = data[key]
            requests.append(pymongo.UpdateOne(filter, {"$set": data}, upsert=True))
        results = coll.bulk_write(requests)
        return results

    def update(self, databaseName, collectionName, target, dataObject):
        db = self.client[databaseName]
        coll = db[collectionName]
        update = {}
        update["$set"] = dataObject
        coll.update_one(target, update, False)



def CollectionTransit(fromDataBase, fromDataBaseName, fromCollectionName, toDataBase, toDatabaseName, toCollectionName):

    #
    fromCollection = fromDataBase.getCollection(fromDataBaseName, fromCollectionName)
    cursor = fromCollection.find()
    #
    #toDataBase.insert_many(toDatabaseName,toCollectionName,cursor)
    #
    #for document in cursor:
    #    del item["_id"]
    #    if item.get("Key") != None:
    #        target = {"Key":item["Key"]}
    #        toDataBase.upsert(toName,toCollection,target,item)
    #    else:
    #        toDataBase.insert(toName,toCollection,item)
    print("Transit Collection Done")


#---Copy Data To Another DB---
#---Filter, 只Copy特定collection_name名字的 collection
def DataBaseTransit(fromDataBase, toDataBase, databaseName, filter = None):

    db = fromDataBase.client[databaseName]
    db2 = toDataBase.client[databaseName]
    coll =  db.collection_names(False)
    process = False
    count = 0
    for collectionName in coll:
        if filter == None:
            process = True
        else:
            if filter in collectionName:
                process = True

        if process == True:
            cursor = db[collectionName].find()
            destCollection = db2[collectionName]
            #
            i = 0
            for document in cursor:
                #i = i + 1
                #print(i)
                del document["_id"]
                wcursor = destCollection.insert_one(document)
            #
            # coll.insert(cursor)#Inseert All
            # Create Index
            destCollection.create_index("Key")
            destCollection.create_index("StdDateTime")
            # CollectionTransit(fromDataBase, databaseName, collectionName, toDataBase, databaseName, collectionName)
            count = count + 1
            print(str(count) + " Copy DataSeries " + collectionName + " " + str(datetime.now()))



def RenameDataBase(database, oldName, newName):

    #db.copyDatabase(oldName, newName);
    db = database.client[oldName]
    coll = db.collection_names()
    #
    db = database.client["admin"]
    for collectionName in coll:
        commandStr = {}
        commandStr["renameCollection"] = oldName + "." + collectionName
        commandStr["to"] = newName + "." + collectionName
        db.command(commandStr)
        print("Renamed :" + collectionName)
        #db.runCommand({"renameCollection": oldName + ".000001.OF_Time_86400_Bar", "to": newName + ".000001.OF_Time_86400_Bar"});
    kkwood = 0